import json
from common.utils import send_http, send_http_stream


class RAG:
    def __init__(self, base_url="http://demo.knowledge.cloudwise.com/knowledgeWeb"):
        self.base_url = base_url
        self.headers = {'accept': 'application/json', 'Content-Type': 'application/json'}
        self.login_headers = {
            'accept': 'application/json',
            'Content-Type': 'application/x-www-form-urlencoded'
        }
        self.login_data = {
            'username': 'admin',
            'password': 'XcZ+s3x6mvC7L6Kze8/O2g/JE58uhFLdh4p35rZJHZB3skaHRZCZz2o='
        }
        self.token = None
        self._login()

    def _login(self):
        """
        调用登录接口获取token
        """
        # 登录时使用form表单格式
        response = send_http(
            self.base_url,
            "/api/v1/token",
            self.login_data,
            self.login_headers
        )
        self.token = response.get('access_token')
        # 更新headers，添加认证信息
        self.headers['Authorization'] = f"Bearer {self.token}"

    def generate(self, content):
        """
        调用流式生成接口
        
        Args:
            content: 用户输入内容
        Returns:
            str: 生成的回复内容
        """
        if not self.token:
            self._login()

        data = {
            "variable": {
                "docs_list": "",
                "message": content
            },
            "session_id": 465,
            "knowledge_id": 30,
            "scene_id": 4
        }

        response = send_http_stream(
            self.base_url,
            "/api/v1/generate_stream",
            data,
            self.headers
        )
        
        for line in response.iter_lines():
            if line:
                raw_data = line.decode('utf-8')
                
                if raw_data.startswith('data:'):
                    try:
                        json_str = raw_data.replace('data:', '').strip()
                        data = json.loads(json_str)
                        
                        if data.get('end') and data.get('loop_symbol'):
                            return data
                    except json.JSONDecodeError:
                        continue
                    except Exception:
                        continue

        return ''
